#!/usr/bin/env python3
#(c) B.Kerler 2021 MIT License
import os
from struct import unpack
from mtkclient.config.payloads import pathconfig
from mtkclient.Library.utils import read_object

entry_region = [
    ('m_buf', 'I'),
    ('m_len', 'I'),
    ('m_start_addr', 'I'),
    ('m_start_offset', 'I'),
    ('m_sig_len', 'I')]

DA = [
    ('magic', 'H'),
    ('hw_code', 'H'),
    ('hw_sub_code', 'H'),
    ('hw_version', 'H'),
    ('sw_version', 'H'),
    ('reserved1', 'H'),
    ('pagesize', 'H'),
    ('reserved3', 'H'),
    ('entry_region_index', 'H'),
    ('entry_region_count', 'H')
    # vector<entry_region> LoadRegion
]

efusedb={}

def main():
    da_setup = []
    loaders = []
    pc = pathconfig()
    for root, dirs, files in os.walk(pc.get_loader_path(), topdown=False):
        for file in files:
            if "MTK_AllInOne" in file:
                loaders.append(os.path.join(root, file))
    if not os.path.exists("loaders"):
        os.mkdir("loaders")
    dadb = {}
    for loader in loaders:
        with open(loader, 'rb') as bootldr:
            bootldr.seek(0x68)
            count_da = unpack("<I", bootldr.read(4))[0]
            for i in range(0, count_da):
                bootldr.seek(0x6C + (i * 0xDC))
                datmp = read_object(bootldr.read(0x14), DA)  # hdr
                da = [datmp]
                # bootldr.seek(0x6C + (i * 0xDC) + 0x14) #sections
                for m in range(0, datmp["entry_region_count"]):
                    entry_tmp = read_object(bootldr.read(20), entry_region)
                    da.append(entry_tmp)
                da_setup.append(da)
                print(f"Loader: {os.path.basename(loader)}")
                dadb[da[0]["hw_code"]]=[("da_payload_addr",da[2]["m_start_addr"]),("pl_payload_addr", da[3]["m_start_addr"])]
                print("hwcode: 0x%04X" % da[0]["hw_code"])
                print("hw_sub_code: 0x%04X" % da[0]["hw_sub_code"])
                print("hw_version: 0x%04X" % da[0]["hw_version"])
                print("sw_version: 0x%04X" % da[0]["sw_version"])
                print("Reserved1: 0x%04X" % da[0]["reserved1"])
                print("Reserved3: 0x%04X" % da[0]["reserved3"])
                for i in range(da[0]["entry_region_count"]):
                    entry=da[i+1]
                    print(f"\t{i}: "+hex(entry['m_start_addr']))
                mbuf=da[3]["m_buf"]
                m_len=da[3]["m_len"]
                startaddr=da[3]["m_start_addr"]
                with open(loader,"rb") as rf:
                    rf.seek(mbuf)
                    data=rf.read(m_len)
                    fname=os.path.join("loaders",hex(da[0]["hw_code"])[2:]+"_"+hex(startaddr)[2:]+os.path.basename(loader))
                    open(fname,"wb").write(data)
                mbuf=da[2]["m_buf"]
                m_len=da[2]["m_len"]
                startaddr=da[2]["m_start_addr"]
                with open(loader,"rb") as rf:
                    rf.seek(mbuf)
                    data=rf.read(m_len)
                    fname=os.path.join("loaders",hex(da[0]["hw_code"])[2:]+"_"+hex(startaddr)[2:]+os.path.basename(loader))
                    open(fname,"wb").write(data)
                print(f"Offset: {hex(mbuf)}")
                print(f"Length: {hex(m_len)}")
                print(f"Addr: {hex(startaddr)}")
                bootldr.seek(da[2]["m_buf"])
                tt=bootldr.read(da[2]["m_len"])
                idx = tt.find(bytes.fromhex("70BB442D27D244A7"))
                #idx = tt.find(bytes.fromhex("01279360D36013615361"))
                if idx!=-1:
                    print("V3 Enabled")
                bootldr.seek(da[3]["m_buf"])
                tt=bootldr.read(da[3]["m_len"])
                idx2 = tt.find(bytes.fromhex("03 29 0D D9 07 4B 1B 68 03 60"))
                if idx2!=-1:
                    efusedb[da[0]["hw_code"]]=hex(int.from_bytes(tt[idx2+0x24:idx2+0x28],'little')&0xFFFFF000)
                else:
                    if not da[0]["hw_code"] in efusedb:
                        efusedb[da[0]["hw_code"]]="None"
                print()

    sorted_dict = dict(sorted(efusedb.items()))
    for hwcode in sorted_dict:
        print(f"[{hex(hwcode)}] efuse_addr = {efusedb[hwcode]}")

    sorted_dict = dict(sorted(dadb.items()))
    for dat in sorted_dict:
        for ldr in dadb[dat]:
            print(f"{hex(dat)}:{ldr[0]}={hex(ldr[1])}")

    """
    chipinfo="/home/bjk/Projects/mtk_bypass/SP_Flash_Tool_v5.2052_Linux/libflashtool.v1.so"
    with open(chipinfo, 'rb') as ci:
        ci.seek(0x3C63C0) #830000004D543635
        data=bytearray(ci.read())
        for i in range(0,len(data),0x48):
            idx=unpack("<I",data[i:i+4])[0]
            name=data[i+4:i+4+0x20].rstrip(b"\x00").decode('utf-8')
            fields=unpack("<HHHIHIIIIII",data[i+4+0x20:i+4+0x20+6+4+2+(6*4)])
            print(str(idx)+" "+name+": "+hex(fields[2]))
    """

if __name__=="__main__":
    main()